from IPython import display
from matplotlib import pyplot as plt
from mxnet.gluon import data as gdata
from mxnet import autograd, gluon, nd
import sys
import time

#下载数据集
mnist_train = gdata.vision.FashionMNIST(train = True)
mnist_test = gdata.vision.FashionMNIST(train = False)

def use_svg_display():
    '''用矢量图显示'''
    display.set_matplotlib_formats('svg')

def set_figsize(figsize):
    '''设置图的尺寸'''
    use_svg_display()
    plt.rcParams['figure.figsize'] = figsize

def draw_scatter(x_vals, y_vals, name = 'f', figsize = (3.5, 2.5)):
    '''画散点图'''
    set_figsize(figsize)
    plt.scatter(x_vals.sum(axis = 1).asnumpy(), y_vals.asnumpy(), 1)
    plt.xlabel('x')
    plt.ylabel(name + '(x)')
    plt.show()

def draw_plot(x_vals, y_vals, name = 'f', figsize = (3.5, 2.5)):
    '''画过点线图'''
    set_figsize(figsize)
    plt.plot(x_vals.asnumpy(), y_vals.asnumpy())
    plt.xlabel('x')
    plt.ylabel(name + '(x)')
    plt.show()

def semilogy(x_vals, y_vals, x_label, y_label, x2_vals = None, y2_vals = None, legend = None, figsize = (3.5, 2.5)):
    '''作图函数_训练误差和泛化误差'''
    set_figsize(figsize)
    plt.xlabel(x_label)
    plt.ylabel(y_label)
    plt.semilogy(x_vals, y_vals)
    if x2_vals and y2_vals:
        plt.semilogy(x2_vals, y2_vals, linestyle = ':')
        plt.legend(legend)
    plt.show()

def get_fashion_mnist_labels(labels):
    '''根据标号转换出名字'''
    text_labels = ['t-shirt', 'trouser', 'pullover', 'dress', 'coat',
                   'sandal', 'shirt', 'sneaker', 'bag', 'ankle boot']
    return [text_labels[int(i)] for i in labels]

def show_fashion_mnist(images, labels):
    '''把图像画出来'''
    use_svg_display()
    _, figs = plt.subplots(1, len(images), figsize = (12, 12))
    for f, img, lbl in zip(figs, images, labels):
        f.imshow(img.reshape((28, 28)).asnumpy())
        f.set_title(lbl)
        f.axes.get_xaxis().set_visible(False)
        f.axes.get_yaxis().set_visible(False)
    plt.show()

def show_net_acc(data, net, num):
    '''显示部分图像验证模型对测试集的准确性'''
    for x, y in data:
        break
    true_labels = get_fashion_mnist_labels(y.asnumpy())
    pred_labels = get_fashion_mnist_labels(net(x).argmax(axis=1).asnumpy())
    titles = [true + '\n' + pred for true, pred in zip(true_labels, pred_labels)]
    show_fashion_mnist(x[0:num], titles[0:num])

def load_data_fashion_mnist(batch_size = 256):
    '''读取数据集'''
    transformer = gdata.vision.transforms.ToTensor()#转换为32位浮点数
    if sys.platform.startswith('win'):
        num_workers = 0
    else:
        num_workers = 4
    train_iter = gdata.DataLoader(mnist_train.transform_first(transformer), batch_size, shuffle = True, num_workers = num_workers)
    test_iter = gdata.DataLoader(mnist_test.transform_first(transformer), batch_size, shuffle = False, num_workers = num_workers)
    return train_iter, test_iter

def evaluate_accuracy(data, net, ctx):
    '''测试模型在数据集上的准确率'''
    acc_sum, n = 0.0, 0
    for x, y in data:
        x = x.as_in_context(ctx)
        y = y.as_in_context(ctx).astype('float32')
        acc_sum += (net(x).argmax(axis = 1) == y).sum().asscalar()
        n += y.size
    return acc_sum / n

def sgd(params, lr, batch_size):
    '''小批量随机梯度下降'''
    for param in params:
        param[ : ] = param - lr * param.grad / batch_size

def train_ch5(net, train_iter, test_iter, num_epochs, ctx, batch_size = 256, params = None, lr = None, trainer = None):
    '''训练模型'''
    loss = gluon.loss.SoftmaxCrossEntropyLoss()
    for epoch in range(num_epochs):
        train_l_sum, train_acc_sum, n = 0.0, 0.0, 0
        for x, y in train_iter:
            x, y = x.as_in_context(ctx), y.as_in_context(ctx)
            with autograd.record():
                y_hat = net(x)
                l = loss(y_hat, y).sum()
            l.backward()
            if trainer is None:
                sgd(params, lr, batch_size)
            else:
                trainer.step(batch_size)
            y = y.astype('float32')
            train_l_sum += l.asscalar()
            train_acc_sum += (y_hat.argmax(axis = 1) == y).sum().asscalar()
            n += y.size
        test_acc = evaluate_accuracy(test_iter, net, ctx)
        print('epoch %d, loss %.4f, train acc %.3f, test acc %.3f' % (epoch + 1, train_l_sum / n, train_acc_sum / n, test_acc))


